#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：exp_analytic_fa.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/28 20:20 
'''
# -*- coding: utf-8 -*-
import os, sys, csv, math, random
from dataclasses import dataclass
from typing import List, Tuple, Dict, Any
import numpy as np
import matplotlib.pyplot as plt
import importlib.util

# ============ CONFIG ============
CORE_FILE = "analytic_fa_core.py"  # set to the filename of YOUR core (the code you pasted)
OUTDIR = "exp_out"
# ================================

def _load_core(core_file: str):
    core_path = os.path.abspath(core_file)
    if not os.path.exists(core_path):
        raise RuntimeError(f"Core file not found: {core_path}. Please set CORE_FILE to your core filename.")
    spec = importlib.util.spec_from_file_location("af_core", core_path)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"Cannot import core module from {core_path}")
    mod = importlib.util.module_from_spec(spec)
    sys.modules["af_core"] = mod
    spec.loader.exec_module(mod)
    return mod, core_path

def ensure_outdir(d=OUTDIR):
    if not os.path.exists(d):
        os.makedirs(d, exist_ok=True)
    return os.path.abspath(d)

def save_csv(path, rows: List[Dict[str,Any]], header: List[str]):
    with open(path, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=header)
        w.writeheader()
        for r in rows: w.writerow(r)

def fronts_to_points(F: np.ndarray, top_k=None):
    if F is None or F.size == 0: return []
    pts = [(-float(f[0]), -float(f[1])) for f in F]
    if top_k:
        pts = sorted(pts, key=lambda x: (-x[0], -x[1]))[:top_k]
    return pts

class FireflyMOX:
    """Wrapper over core FireflyMO with knobs: init_mix and use_SA."""
    def __init__(self, core, plant, prm, fap, sharing_on=True, init_mix=(1/3,1/3,1/3), use_SA=True):
        self.core = core
        self.inner = core.FireflyMO(plant, prm, fap, sharing_on)
        self.init_mix = init_mix
        self.use_SA = use_SA
        # expose needed refs
        self.plant, self.prm, self.fap, self.sharing_on = plant, prm, fap, sharing_on
        self.Cmin = self.inner.Cmin; self.Ts_allow = self.inner.Ts_allow; self._scale_C_total = self.inner._scale_C_total
        self.pack = self.inner.pack; self.unpack = self.inner.unpack; self.fitness = self.inner.fitness
        self.init_random = self.inner.init_random; self.init_gauss = self.inner.init_gauss; self.init_rate_based = self.inner.init_rate_based
        self.refine_Ts_SA = self.inner.refine_Ts_SA

    def run(self, verbose=True):
        P, G = self.fap.pop, self.fap.gens
        w_r, w_g, w_t = self.init_mix
        n_r = int(round(P * w_r)); n_g = int(round(P * w_g)); n_t = P - n_r - n_g
        inits = []
        for _ in range(max(0, n_r)): inits.append(self.init_random())
        for _ in range(max(0, n_g)): inits.append(self.init_gauss())
        for _ in range(max(0, n_t)): inits.append(self.init_rate_based())
        while len(inits) < P: inits.append(self.init_random())
        X = np.stack([self.pack(C, Ts) for C,Ts in inits], axis=0)

        F = np.zeros((P,2)); MET=[None]*P
        for i in range(P): F[i], MET[i] = self.fitness(X[i])

        archX=[]; archF=[]
        def update_archive(Xlist, Flist):
            nonlocal archX, archF
            allX = archX + [x.copy() for x in Xlist]
            allF = archF + [f.copy() for f in Flist]
            Farr = np.stack(allF, axis=0)
            fronts = self.core.fast_non_dominated_sort(Farr)
            newX, newF = [], []
            for idx in fronts[0]:
                newX.append(allX[idx].copy()); newF.append(allF[idx].copy())
            archX, archF = newX, newF

        update_archive(list(X), list(F))

        for g in range(1, G+1):
            fronts = self.core.fast_non_dominated_sort(F)
            elite_idx = []
            for fr in fronts:
                if len(elite_idx) >= self.fap.elite: break
                cd = self.core.crowding_distance(fr, F); order = np.argsort(-cd)
                for k in order:
                    if len(elite_idx) < self.fap.elite:
                        elite_idx.append(fr[k])

            elites = X[elite_idx].copy()
            attr0, gamma = self.fap.attr0, self.fap.gamma
            kappa = self.fap.kappa_hi - (self.fap.kappa_hi - self.fap.kappa_lo) * (g / G)

            Y = []
            while len(Y) < P - self.fap.elite:
                a, b = np.random.randint(0, P), np.random.randint(0, P)
                xa, xb = X[a].copy(), X[b].copy()
                Fa, Fb = F[a], F[b]
                move_from, move_to = (xa, xb) if np.any(Fa > Fb) else (xb, xa)
                Ca, Ta = self.unpack(move_from)
                Cb, Tb = self.unpack(move_to)
                d = np.linalg.norm(move_from - move_to)
                attr = attr0 * math.exp(-gamma * (d**2))
                Cnew = Ca + attr*(Cb - Ca) + np.random.randn(*Ca.shape)*kappa*10.0
                Tsnew = Ta + attr*(Tb - Ta) + np.random.randn(*Ta.shape)*kappa
                Cnew = np.maximum(self.Cmin, np.round(Cnew)).astype(int)
                self._scale_C_total(Cnew)
                Tsnew = np.clip(Tsnew, 0.0, 1.0) * self.Ts_allow
                if self.use_SA:
                    Tsnew = self.refine_Ts_SA(Cnew, Tsnew)
                Y.append(self.pack(Cnew, Tsnew))

            Xnew = np.vstack([elites] + Y)
            Fnew = np.zeros_like(F); METnew=[None]*P
            for i in range(P): Fnew[i], METnew[i] = self.fitness(Xnew[i])

            Xall = np.vstack([X, Xnew]); Fall = np.vstack([F, Fnew])
            fronts_all = self.core.fast_non_dominated_sort(Fall)
            Z=[]
            for fr in fronts_all:
                if len(Z)+len(fr) <= P:
                    Z += fr
                else:
                    cd = self.core.crowding_distance(fr, Fall)
                    need = P - len(Z)
                    order = np.argsort(-cd)[:need]
                    Z += [fr[k] for k in order]
                    break
            X = Xall[Z]; F = Fall[Z]
            update_archive(list(X), list(F))

            if verbose and (g % 20 == 0 or g == 1):
                bestW = -np.min(F[:,0]); bestTh = -np.min(F[:,1])
                print(f"[Gen {g:4d} | {'ON' if self.inner.sharing_on else 'OFF'}] front={len(archX)}  best W≈{bestW:.2f}  θ≈{bestTh:.4f}")

        details=[]
        for x in archX:
            C, Ts = self.unpack(x)
            _, met = self.fitness(self.pack(C, Ts))
            details.append((C, Ts, met))
        return np.stack(archF, axis=0), details

def scale_C_total(prm, new_total: int):
    return af_core.EvalParams(C_max_total=new_total, C_min_each=prm.C_min_each, time_horizon=prm.time_horizon,
                              transport_fee_scale=prm.transport_fee_scale, w_negW=prm.w_negW, w_negTheta=prm.w_negTheta)

def scale_transport_fee(prm, factor: float):
    return af_core.EvalParams(C_max_total=prm.C_max_total, C_min_each=prm.C_min_each, time_horizon=prm.time_horizon,
                              transport_fee_scale=prm.transport_fee_scale * factor, w_negW=prm.w_negW, w_negTheta=prm.w_negTheta)

def scale_reliability(plant, mtbf_scale=1.0, mttr_scale=1.0):
    lines = []
    for ln in plant.lines:
        ms = [af_core.MachineParam(m.omega, m.MTBF_h*mtbf_scale, m.MTTR_h*mttr_scale) for m in ln.machines]
        lines.append(af_core.Line(machines=ms, unit_earnings=ln.unit_earnings[:], input_qty=ln.input_qty))
    return af_core.Plant(lines=lines, share_mask=plant.share_mask.copy(), transport_cost=plant.transport_cost.copy())

def skew_revenue(plant, skew_factor=1.4):
    lines = []
    for ln in plant.lines:
        vals = ln.unit_earnings[:]
        mx = max(vals); idx = vals.index(mx)
        new_vals = vals[:]
        new_vals[idx] = mx * skew_factor
        s_old, s_new = sum(vals), sum(new_vals)
        if s_new > 0:
            new_vals = [v * s_old / s_new for v in new_vals]
        lines.append(af_core.Line(machines=ln.machines[:], unit_earnings=new_vals, input_qty=ln.input_qty))
    return af_core.Plant(lines=lines, share_mask=plant.share_mask.copy(), transport_cost=plant.transport_cost.copy())

def exp1_baselines(core, outdir, seed=1):
    print("\n=== Exp1: Baselines (Sharing ON vs OFF) ===")
    random.seed(seed); np.random.seed(seed)
    plant, prm = core.build_demo_plant(n_lines=3, m_machines=9, seed=7)
    cfg = core.FAParams(pop=60, gens=120, elite=6, seed=seed)
    # ON
    f_on = FireflyMOX(core, plant, prm, cfg, sharing_on=True, init_mix=(1/3,1/3,1/3), use_SA=True)
    F_on, _ = f_on.run(verbose=True)
    # OFF
    f_off = FireflyMOX(core, plant, prm, cfg, sharing_on=False, init_mix=(1/3,1/3,1/3), use_SA=True)
    F_off, _ = f_off.run(verbose=True)

    pts_on  = fronts_to_points(F_on, top_k=30)
    pts_off = fronts_to_points(F_off, top_k=30)

    rows = []
    for mode, pts in [("ON", pts_on), ("OFF", pts_off)]:
        for i,(W,th) in enumerate(pts, 1):
            rows.append({"mode": mode, "rank": i, "W": f"{W:.4f}", "theta": f"{th:.6f}"})
    save_csv(os.path.join(outdir, "exp1_baselines.csv"), rows, ["mode","rank","W","theta"])

    plt.figure()
    if pts_on:
        th_on = [t for _,t in pts_on]; W_on = [w for w,_ in pts_on]
        plt.scatter(th_on, W_on, label="Sharing ON")
    if pts_off:
        th_off = [t for _,t in pts_off]; W_off = [w for w,_ in pts_off]
        plt.scatter(th_off, W_off, label="Sharing OFF")
    plt.xlabel("Mean throughput θ"); plt.ylabel("Total earnings W"); plt.title("Exp1: Pareto Fronts")
    plt.legend(); plt.tight_layout(); plt.savefig(os.path.join(outdir, "exp1_fronts.png")); plt.close()

def exp2_ablations(core, outdir, seed=1):
    print("\n=== Exp2: Ablations (init mix / SA / multi-start) ===")
    random.seed(seed); np.random.seed(seed)
    plant, prm = core.build_demo_plant(n_lines=3, m_machines=9, seed=7)
    base = core.FAParams(pop=50, gens=100, elite=6, seed=seed)
    cases = [
        ("full",          dict(init_mix=(1/3,1/3,1/3), use_SA=True,  multi_start=1)),
        ("no_SA",         dict(init_mix=(1/3,1/3,1/3), use_SA=False, multi_start=1)),
        ("only_random",   dict(init_mix=(1.0,0.0,0.0), use_SA=True,  multi_start=1)),
        ("only_gauss",    dict(init_mix=(0.0,1.0,0.0), use_SA=True,  multi_start=1)),
        ("only_rate",     dict(init_mix=(0.0,0.0,1.0), use_SA=True,  multi_start=1)),
        ("multi_start3",  dict(init_mix=(1/3,1/3,1/3), use_SA=True,  multi_start=3)),
    ]
    rows = []
    for name, kw in cases:
        archives = []
        ms = kw["multi_start"]
        for s in range(ms):
            cfg = core.FAParams(pop=base.pop, gens=base.gens, elite=base.elite, seed=seed+s)
            solver = FireflyMOX(core, plant, prm, cfg, sharing_on=True, init_mix=kw["init_mix"], use_SA=kw["use_SA"])
            F, _ = solver.run(verbose=False)
            archives.append(F)
        allF = np.vstack(archives)
        fronts = core.fast_non_dominated_sort(allF)
        first = allF[fronts[0]]
        pts = fronts_to_points(first, top_k=20)
        for i,(W,th) in enumerate(pts,1):
            rows.append({"case": name, "rank": i, "W": f"{W:.4f}", "theta": f"{th:.6f}"})
        # plot
        plt.figure()
        if pts:
            th=[t for _,t in pts]; W=[w for w,_ in pts]
            plt.scatter(th, W)
        plt.xlabel("θ"); plt.ylabel("W"); plt.title(f"Exp2: {name}")
        plt.tight_layout(); plt.savefig(os.path.join(outdir, f"exp2_{name}.png")); plt.close()
    save_csv(os.path.join(outdir, "exp2_ablations.csv"), rows, ["case","rank","W","theta"])

def exp3_sensitivity(core, outdir, seed=1):
    print("\n=== Exp3: Sensitivity ===")
    random.seed(seed); np.random.seed(seed)
    plant, prm = core.build_demo_plant(n_lines=3, m_machines=9, seed=7)
    base = core.FAParams(pop=50, gens=100, elite=6, seed=seed)
    rows = []
    # (a) C_total sweep
    for C in (800, 1200, 1600, 2000):
        prmC = scale_C_total(prm, C)
        solver = FireflyMOX(core, plant, prmC, base, sharing_on=True)
        F, _ = solver.run(verbose=False)
        if F.size>0:
            W, th = -float(F[:,0].min()), -float(F[:,1].min())
            rows.append({"type":"C_total","level":C,"W":f"{W:.4f}","theta":f"{th:.6f}"})
    # (b) transport fee factor
    for fac in (0.5, 1.0, 2.0, 3.0):
        prmT = scale_transport_fee(prm, fac)
        solver = FireflyMOX(core, plant, prmT, base, sharing_on=True)
        F, _ = solver.run(verbose=False)
        if F.size>0:
            W, th = -float(F[:,0].min()), -float(F[:,1].min())
            rows.append({"type":"TransportFee","level":fac,"W":f"{W:.4f}","theta":f"{th:.6f}"})
    # (c) reliability scaling
    for mtbf_s, mttr_s in ((1.2,0.8),(1.0,1.0),(0.8,1.2)):
        plantR = scale_reliability(plant, mtbf_s, mttr_s)
        solver = FireflyMOX(core, plantR, prm, base, sharing_on=True)
        F, _ = solver.run(verbose=False)
        if F.size>0:
            W, th = -float(F[:,0].min()), -float(F[:,1].min())
            rows.append({"type":"Reliability","level":f"MTBFx{mtbf_s}_MTTRx{mttr_s}","W":f"{W:.4f}","theta":f"{th:.6f}"})
    # (d) revenue skew
    for k in (1.0, 1.3, 1.6):
        plantS = skew_revenue(plant, k)
        solver = FireflyMOX(core, plantS, prm, base, sharing_on=True)
        F, _ = solver.run(verbose=False)
        if F.size>0:
            W, th = -float(F[:,0].min()), -float(F[:,1].min())
            rows.append({"type":"RevenueSkew","level":k,"W":f"{W:.4f}","theta":f"{th:.6f}"})
    save_csv(os.path.join(outdir, "exp3_sensitivity.csv"), rows, ["type","level","W","theta"])
    ct = [r for r in rows if r["type"]=="C_total"]
    if ct:
        X = [int(r["level"]) for r in ct]
        W = [float(r["W"]) for r in ct]
        th= [float(r["theta"]) for r in ct]
        plt.figure(); plt.plot(X, W, marker='o'); plt.xlabel("C_total"); plt.ylabel("W"); plt.title("Exp3: W vs C_total"); plt.tight_layout(); plt.savefig(os.path.join(outdir,"exp3_W_vs_C.png")); plt.close()
        plt.figure(); plt.plot(X, th, marker='o'); plt.xlabel("C_total"); plt.ylabel("θ"); plt.title("Exp3: θ vs C_total"); plt.tight_layout(); plt.savefig(os.path.join(outdir,"exp3_theta_vs_C.png")); plt.close()

def exp4_scalability(core, outdir, seed=1):
    print("\n=== Exp4: Scalability ===")
    random.seed(seed); np.random.seed(seed)
    settings = [(3,9),(4,10),(6,12)]
    rows=[]
    for nL, nM in settings:
        plant, prm = core.build_demo_plant(n_lines=nL, m_machines=nM, seed=7)
        cfg = core.FAParams(pop=70, gens=120, elite=6, seed=seed)
        solver = FireflyMOX(core, plant, prm, cfg, sharing_on=True)
        F, _ = solver.run(verbose=False)
        pts = fronts_to_points(F, top_k=30)
        for i,(W,th) in enumerate(pts,1):
            rows.append({"lines":nL,"stages":nM,"rank":i,"W":f"{W:.4f}","theta":f"{th:.6f}"})
        plt.figure()
        if pts:
            ths=[t for _,t in pts]; Ws=[w for w,_ in pts]
            plt.scatter(ths, Ws)
        plt.xlabel("θ"); plt.ylabel("W"); plt.title(f"Exp4: Pareto (lines={nL}, stages={nM})")
        plt.tight_layout(); plt.savefig(os.path.join(outdir, f"exp4_{nL}x{nM}.png")); plt.close()
    save_csv(os.path.join(outdir, "exp4_scalability.csv"), rows, ["lines","stages","rank","W","theta"])

def main():
    core, core_path = _load_core(CORE_FILE)
    outdir = ensure_outdir()
    print("Loaded core from:", core_path)
    exp1_baselines(core, outdir, seed=1)
    exp2_ablations(core, outdir, seed=1)
    exp3_sensitivity(core, outdir, seed=1)
    exp4_scalability(core, outdir, seed=1)
    print("All experiments completed. Results saved to:", outdir)

if __name__ == "__main__":
    main()
